import cv2
import numpy as np
import math
import os





def takeX(L):
    return  L[0][0] 

def get_ellipse_param(major_radius, minor_radius, angle):
    a, b = major_radius, minor_radius
    sin_theta = np.sin(-angle)
    cos_theta = np.cos(-angle)
    A = a**2 * sin_theta**2 + b**2 * cos_theta**2
    B = 2 * (a**2 - b**2) * sin_theta * cos_theta
    C = a**2 * cos_theta**2 + b**2 * sin_theta**2
    F = -a**2 * b**2
    return A, B, C, F


def calculate_rectangle(A, B, C, F):
    '''
    椭圆上下外接点的纵坐标值
    '''
    y = np.sqrt(4*A*F / (B**2 - 4*A*C))
    y1, y2 = -np.abs(y), np.abs(y)
    
    '''
    椭圆左右外接点的横坐标值
    '''
    x = np.sqrt(4*C*F / (B**2 - 4*C*A))
    x1, x2 = -np.abs(x), np.abs(x)
    
    return (x1, y1), (x2, y2)

def get_rectangle(major_radius, minor_radius, angle, center_x, center_y):
    A, B, C, F = get_ellipse_param(major_radius, minor_radius, angle)
    p1, p2 = calculate_rectangle(A, B, C, F)
    return (center_x+p1[0], center_y+p1[1]), (center_x+p2[0], center_y+p2[1])







def get_boards(path,num):
    t_num = num
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    (b, _, _) = cv2.split(img)
    _, scr = cv2.threshold(b, 100, 255, cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,5)) #腐蚀、膨胀核
    opened = cv2.morphologyEx(scr, cv2.MORPH_OPEN, kernel) #开操作
    closed = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, kernel) #闭操作

    lim_dir_area = 10
    lim_pro_area = 0.6
    lim_pro = 0.2

    lights = []

    #temp = np.ones(img.shape,np.uint8)*255
    contours, hierarchy=cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) #边缘检测
    for contour in contours:
        lightArea = cv2.contourArea(contour)
        if lightArea <= lim_dir_area:
            continue

        lightRec = cv2.fitEllipse(contour)
        RecArea = lightRec[1][0] * lightRec[1][1]

        if lightArea / RecArea <= lim_pro_area:
            continue

        if lightRec[1][0]/lightRec[1][1] >= lim_pro:
            continue

        lights.append(lightRec)
        #cv2.ellipse(temp,lightRec,(0,255,0),1)

    lights.sort(key=takeX)

    l = len(lights)
    lim_angle = 2
    lim_temp_angle = 30


    boards = []
    for i in range(l):
        for j in range(i+1,l):
            dAngle = abs(lights[i][2] - lights[j][2])
            dx = lights[j][0][0] - lights[i][0][0]
            dy = lights[j][0][1] - lights[i][0][1]



            temp = math.atan(dy/dx) * 180/math.pi
            if lights[i][2] > 90:
                dif_angle = abs(lights[i][2] - 180 + temp)
            else:
                dif_angle = abs(lights[i][2] + temp)

            if dAngle < lim_angle and dif_angle < lim_temp_angle:
                boards.append((i,j))
                break
    
    #croped_region = img
    for board in boards:
        a = lights[board[0]]
        b = lights[board[1]]

        p1,p2 = get_rectangle(a[1][0],a[1][1],a[2] * np.pi /180 ,a[0][0],a[0][1])
        p3,p4 = get_rectangle(b[1][0],b[1][1],b[2] * np.pi /180 ,b[0][0],b[0][1])

        min_x = int(min(p1[0],p2[0],p3[0],p4[0]))
        min_y = int(min(p1[1],p2[1],p3[1],p4[1]))
        max_x = int(max(p1[0],p2[0],p3[0],p4[0]))
        max_y = int(max(p1[1],p2[1],p3[1],p4[1]))
        #print(min_x,min_y,max_x,max_y)
        croped_region = img[min_y:max_y,min_x:max_x]
        croped = cv2.resize(croped_region,(64,64))
        #cv2.rectangle(img,(min_x, max_y),(max_x, min_y),(255,0,0),1)
        cv2.imwrite('./data/1/4/%d.jpg'%t_num,croped)
        t_num = t_num + 1
        cv2.imshow("test", croped)
        cv2.waitKey(33)

    return t_num


path = './img/02'
num = 601
file_list = os.listdir(path)
file_list.sort()
for file in range(708):
    file_path = os.path.join(path,"%d.jpg"%file)
    print(file_path)
    num = get_boards(file_path,num)

cv2.waitKey(0)
cv2.destroyAllWindows()